import gdal
import numpy as np
from osgeo import gdal
from tqdm import tqdm  # 用于显示进度条

from utils import find_files


class GeoTiff:
    def __init__(self, file_path):
        """
        初始化GeoTiff对象

        参数:
        file_path -- GeoTIFF文件的路径
        """
        self.file_path = file_path
        self.dataset = gdal.Open(file_path)

        if not self.dataset:
            raise FileNotFoundError(f"无法打开文件: {file_path}")

        # 获取地理参考信息
        self.projection = self.dataset.GetProjection()
        self.geo_transform = self.dataset.GetGeoTransform()
        self.cols = self.dataset.RasterXSize
        self.rows = self.dataset.RasterYSize

        # 读取第一个波段的数据
        band1 = self.dataset.GetRasterBand(1)
        self.band1_data = np.array(band1.ReadAsArray())
        self.nodata_value = band1.GetNoDataValue()  # 保存原始nodata值

    def __del__(self):
        """析构函数，释放资源"""
        self.dataset = None

    def get_value(self, col, row):
        return self.band1_data[col, row]

    def get_latlon(self, m, n):
        """
        根据像素位置获取经纬度
        :param m: 水平像素位置
        :param n: 垂直像素位置
        :return: 对应的经纬度 (lat, lon)
        """
        lon = self.geo_transform[0] + m * self.geo_transform[1] + n * self.geo_transform[2]
        lat = self.geo_transform[3] + m * self.geo_transform[4] + n * self.geo_transform[5]
        return lat, lon

    def get_pixel_value(self, lat, lon):
        """
        根据经纬度获取像素值
        :param lat: 纬度
        :param lon: 经度
        :return: 对应像素值
        """
        inv_geo_transform = gdal.InvGeoTransform(self.geo_transform)
        x, y = gdal.ApplyGeoTransform(inv_geo_transform, lon, lat)
        x, y = int(x), int(y)

        if x < 0 or x >= self.dataset.RasterXSize or y < 0 or y >= self.dataset.RasterYSize:
            return np.nan

        return self.band1_data[y, x]

    def get_pixel_value_interpolated(self, lat, lon):
        """
        根据经纬度获取插值后的像素值
        :param lat: 纬度
        :param lon: 经度
        :return: 插值后的像素值
        """
        inv_geo_transform = gdal.InvGeoTransform(self.geo_transform)
        x, y = gdal.ApplyGeoTransform(inv_geo_transform, lon, lat)

        x0, y0 = int(np.floor(x)), int(np.floor(y))
        x1, y1 = x0 + 1, y0 + 1

        # 边界检查
        if x0 < 0 or x1 >= self.dataset.RasterXSize or y0 < 0 or y1 >= self.dataset.RasterYSize:
            return np.nan

        # 进行双线性插值
        f00 = self.band1_data[y0, x0]
        f01 = self.band1_data[y1, x0]
        f10 = self.band1_data[y0, x1]
        f11 = self.band1_data[y1, x1]

        if np.isnan(f00) or np.isnan(f01) or np.isnan(f10) or np.isnan(f11):
            return np.nan

        fxy1 = f00 * (x1 - x) + f10 * (x - x0)
        fxy2 = f01 * (x1 - x) + f11 * (x - x0)
        fxy = fxy1 * (y1 - y) + fxy2 * (y - y0)

        return fxy

    def save_as(self, save_path: str, data: np.ndarray = None, nodata_value: float = None):
        """
        将数据保存为GeoTIFF文件

        参数:
        data -- 要保存的二维数组
        save_path -- 保存路径
        nodata_value -- 无效值标记
        """
        if data is None:
            data = self.band1_data

        # 创建输出文件
        driver = gdal.GetDriverByName("GTiff")
        out_dataset = driver.Create(
            save_path,
            self.cols,
            self.rows,
            1,  # 波段数量
            gdal.GDT_Float32  # 数据类型
        )

        # 设置地理参考信息
        out_dataset.SetProjection(self.projection)
        out_dataset.SetGeoTransform(self.geo_transform)

        # 写入数据
        out_band = out_dataset.GetRasterBand(1)
        out_band.WriteArray(data)

        # 设置无效值
        if nodata_value is not None:
            out_band.SetNoDataValue(nodata_value)

        # 清理资源
        out_band.FlushCache()
        out_dataset = None


class GeoTiffBatch:

    def __init__(self, tiff_dir=None, tiff_paths=None):
        file_paths = tiff_paths
        if not file_paths:
            if tiff_dir:
                file_paths = find_files(tiff_dir, lambda f: f.upper().endswith('.TIF'))
        if not file_paths:
            raise ValueError("empty list/dir")
        self.geo_tiffs = []
        for file_path in file_paths:
            geo_tiff = GeoTiff(file_path)
            self.geo_tiffs.append(geo_tiff)
        self.band1_data = np.array([geo_tiff.band1_data for geo_tiff in self.geo_tiffs])
        pass

    def max_at_mn(self, m, n):
        return np.nanmax(self.band1_data[:, m, n])

    def mean_at_mn(self, m, n):
        return np.nanmean(self.band1_data[:, m, n])

    def do_overlap(self, check_func, save_path=None):
        """
        合并多个GeoTIFF文件，优先使用有效值填充无效区域

        参数:
        check_func -- 检查值是否有效的函数
        save_path -- 合并后的TIFF文件保存路径
        """
        if not self.geo_tiffs:
            raise ValueError("没有可用的GeoTIFF文件")

        # 获取基准数据
        base_tiff = self.geo_tiffs[0]
        base_data = base_tiff.band1_data.copy()

        # 创建初始掩码
        mask = ~np.vectorize(check_func)(base_data)

        # 如果有多个TIFF，尝试填充
        if len(self.geo_tiffs) > 1:
            print(f"开始合并 {len(self.geo_tiffs)} 个文件...")

            # 如果有多个TIFF，尝试用后续文件填充无效区域
            if len(self.geo_tiffs) > 1:
                print(f"开始合并 {len(self.geo_tiffs)} 个文件...")

                for idx, tiff in enumerate(tqdm(self.geo_tiffs[1:], desc="合并文件中")):
                    current_data = tiff.band1_data

                    # 找出需要替换的位置：掩码为真且当前值有效
                    valid_mask = np.vectorize(check_func)(current_data)
                    replace_mask = mask & valid_mask

                    # 替换值并更新掩码
                    base_data[replace_mask] = current_data[replace_mask]
                    mask = mask & ~valid_mask

                    # 如果所有无效像素都已填充，提前结束
                    if not np.any(mask):
                        print(f"所有无效像素已在第 {idx + 1} 个文件后填充完毕")
                        break

        if save_path:
            # 保存结果
            base_tiff.save_as(save_path, base_data)
            print(f"合并文件已保存至: {save_path}")

        return base_data

    def do_aggregate(self, check_func, aggregate_func, save_path=None):
        """
        聚合多个GeoTIFF文件，对每个像素位置应用聚合函数

        参数:
        check_func -- 检查值是否有效的函数
        aggregate_func -- 聚合函数（如np.nanmax, np.nanmean等）
        save_path -- 保存路径
        nodata_value -- 输出文件中使用的无效值
        """
        if not self.geo_tiffs:
            raise ValueError("没有可用的GeoTIFF文件")

        # 获取第一个TIFF的形状作为基准
        base_tiff = self.geo_tiffs[0]
        base_shape = base_tiff.band1_data.shape

        # 创建结果数组
        result_data = np.full(base_shape, base_tiff.nodata_value, dtype=np.float32)

        # 创建三维掩码：标记所有有效值位置
        valid_mask = np.vectorize(check_func)(self.band1_data)

        # 对每个像素位置进行处理
        print("开始聚合操作...")
        for i in tqdm(range(base_shape[0]), desc="处理行"):
            for j in range(base_shape[1]):
                # 获取该位置所有有效值
                valid_values = self.band1_data[:, i, j][valid_mask[:, i, j]]

                # 如果有有效值，应用聚合函数
                if valid_values.size > 0:
                    result_data[i, j] = aggregate_func(valid_values)

        if save_path:
            # 保存结果
            base_tiff.save_as(save_path, result_data)
            print(f"聚合文件已保存至: {save_path}")

        return result_data

    def close(self):
        """显式释放所有资源"""
        for tiff in self.geo_tiffs:
            tiff.__del__()
        self.geo_tiffs = []
